#include "TabTipSystemBase.h"
#include "wndipc\IPCWndServer.h"
#include "wndipc\IPCWndClient.h"
#include <stdint.h>
#include <assert.h>
#include <Windows.h>

#define ImmersiveFocusEventName L"ImmersiveFocusTrackingActiveEvent"
#define WM_SHOWINPUTDBGWND (WM_USER + 0x111)
#define WM_HIDEINPUTDBGWND (WM_USER + 0x112)
#define g_info printf
#define g_warning printf

typedef struct _WNF_TYPE_ID {
    GUID         TypeId;
} WNF_TYPE_ID, * PWNF_TYPE_ID;

typedef const WNF_TYPE_ID* PCWNF_TYPE_ID;

typedef struct _WNF_CONTEXT_HEADER {
    USHORT                            NodeTypeCode;
    USHORT                            NodeByteSize;
} WNF_CONTEXT_HEADER, * PWNF_CONTEXT_HEADER;

typedef struct _WNF_STATE_NAME_INTERNAL {
    ULONG64 Version : 4;
    ULONG64 NameLifetime : 2;
    ULONG64 DataScope : 4;
    ULONG64 PermanentData : 1;
    ULONG64 Unique : 53;
} WNF_STATE_NAME_INTERNAL, * PWNF_STATE_NAME_INTERNAL;

typedef ULONG WNF_CHANGE_STAMP, * PWNF_CHANGE_STAMP;

typedef struct _WNF_STATE_NAME {
    ULONG                             Data[2];
} WNF_STATE_NAME, * PWNF_STATE_NAME;

typedef struct _WNF_DELIVERY_DESCRIPTOR {
    ULONG64                           SubscriptionId;
    WNF_STATE_NAME                    StateName;
    WNF_CHANGE_STAMP                  ChangeStamp;
    ULONG                             StateDataSize;
    ULONG                             EventMask;
    WNF_TYPE_ID                       TypeId;
    ULONG                             StateDataOffset;
} WNF_DELIVERY_DESCRIPTOR, * PWNF_DELIVERY_DESCRIPTOR;

typedef struct _WNF_NAME_SUBSCRIPTION {
    WNF_CONTEXT_HEADER                Header;
    ULONG64                           SubscriptionId;
    WNF_STATE_NAME_INTERNAL           StateName;
    WNF_CHANGE_STAMP                  CurrentChangeStamp;
    LIST_ENTRY                        NamesTableEntry;
    PWNF_TYPE_ID                      TypeId;
    SRWLOCK                           SubscriptionLock;
    LIST_ENTRY                        SubscriptionsListHead;
    ULONG                             NormalDeliverySubscriptions;
    ULONG                             NotificationTypeCount[5];
    PWNF_DELIVERY_DESCRIPTOR          RetryDescriptor;
    ULONG                             DeliveryState;
    ULONG64                           ReliableRetryTime;
} WNF_NAME_SUBSCRIPTION, * PWNF_NAME_SUBSCRIPTION;

typedef struct _WNF_SERIALIZATION_GROUP {
    WNF_CONTEXT_HEADER                Header;
    ULONG                             GroupId;
    LIST_ENTRY                        SerializationGroupList;
    ULONG64                           SerializationGroupValue;
    ULONG64                           SerializationGroupMemberCount;
} WNF_SERIALIZATION_GROUP, * PWNF_SERIALIZATION_GROUP;

typedef NTSTATUS(*PWNF_USER_CALLBACK) (
    WNF_STATE_NAME                    StateName,
    WNF_CHANGE_STAMP                  ChangeStamp,
    PWNF_TYPE_ID                      TypeId,
    PVOID                             CallbackContext,
    PVOID                             Buffer,
    ULONG                             BufferSize);

typedef struct _WNF_USER_SUBSCRIPTION
{
    WNF_CONTEXT_HEADER Header;
    LIST_ENTRY SubscriptionsListEntry;
    PWNF_NAME_SUBSCRIPTION NameSubscription;
    PWNF_USER_CALLBACK Callback;
    PVOID CallbackContext;
    ULONG64 SubProcessTag;
    ULONG CurrentChangeStamp;
    ULONG DeliveryOptions;
    ULONG SubscribedEventSet;
    PWNF_SERIALIZATION_GROUP SerializationGroup;
    ULONG UserSubscriptionCount;
    ULONG64 Unknown[10];
}
WNF_USER_SUBSCRIPTION, * PWNF_USER_SUBSCRIPTION;

static WNF_STATE_NAME WNF_TKBN_IMMERSIVE_FOCUS_TRACKING{ 0xA3BC1835, 0x0F840539 };
static PWNF_USER_SUBSCRIPTION g_pStaUserSubscription = NULL;

HWND    CTabTipSystemBase::m_hStaFocusEdit = NULL;
HWND    CTabTipSystemBase::m_hStaFocusAncestor = NULL;
HWND    CTabTipSystemBase::m_hStaFocusEditPointHWnd = NULL;
bool    CTabTipSystemBase::m_bStaStartListenMouseUp = false;
bool    CTabTipSystemBase::m_bStaShowingInputWnd = false;
RECT    CTabTipSystemBase::m_rtStaFocusEdit = { 0,0,0,0 };
CIPCWndClient* CTabTipSystemBase::ms_pIPCWndClient = NULL;

CTabTipSystemBase::CTabTipSystemBase() 
    : m_hImmersiveEvent(NULL)
    , m_hTabtipNotifyDbgWnd(NULL)
    , m_pIPCWndServer(NULL)
{
    memset(m_szBufImmersiveFocusEventName, 0, sizeof(m_szBufImmersiveFocusEventName));
    ObtainDbgWndHWnd();

    m_pIPCWndServer = new CIPCWndServer;
    if (!CTabTipSystemBase::ms_pIPCWndClient)
    {
        CTabTipSystemBase::ms_pIPCWndClient = new CIPCWndClient;
    }
}

CTabTipSystemBase::~CTabTipSystemBase()
{
    SAFEDELEPTR(m_pIPCWndServer);
    SAFEDELEPTR(CTabTipSystemBase::ms_pIPCWndClient);
}

BOOL CTabTipSystemBase::InstallNotification()
{
    Initialize();
    InstallWndForIPC();
    InstallImmersiveFocusTracking();
    return TRUE;
}

BOOL CTabTipSystemBase::UninstallNotification()
{
    UninstallImmersiveFocusTracking();
    Uninitialize();
    UninstallWndForIPC();
    return TRUE;
}

BOOL CTabTipSystemBase::InstallWndForIPC()
{
    m_pIPCWndServer->set_server_name(L"TabTipNotificationIPCSVER");
    m_pIPCWndServer->set_handler(this); // 必须先设置回调  然后再listen
    m_pIPCWndServer->start_listen();

    ms_pIPCWndClient->set_host_server_hwnd(m_pIPCWndServer->get_hwnd());
    return TRUE;
}

BOOL CTabTipSystemBase::UninstallWndForIPC()
{
    ms_pIPCWndClient->set_host_server_hwnd(NULL);
    m_pIPCWndServer->stop_listen();
    return TRUE;
}

void CTabTipSystemBase::ObtainDbgWndHWnd()
{
    HANDLE hShareMemoryFileMap = OpenFileMappingW(FILE_MAP_READ, 0, L"TabTipNotificationDbgWndShareMemory");
    if (hShareMemoryFileMap)
    {
        LPVOID lpViewFile = MapViewOfFile(hShareMemoryFileMap, FILE_MAP_READ, 0, 0, 0);
        if (lpViewFile)
        {
            memcpy(&m_hTabtipNotifyDbgWnd, lpViewFile, sizeof(HWND));
        }
        CloseHandle(hShareMemoryFileMap);
    }
}

LRESULT CTabTipSystemBase::OnProcessIPCMsg(UINT msg, WPARAM wParam, LPARAM lParam)
{
    if (msg == WM_TABTIP_DISPLAYINPUTWND)
    {
        if (wParam == SW_SHOW)
        {
            if (CTabTipSystemBase::m_bStaShowingInputWnd)
            {
                return S_OK;
            }
            CTabTipSystemBase::m_bStaShowingInputWnd = true;

            if (m_hTabtipNotifyDbgWnd)
            {
                ::PostMessage(m_hTabtipNotifyDbgWnd, WM_SHOWINPUTDBGWND, 0, 0);
            }
        }
        else
        {  // hide
            if (!CTabTipSystemBase::m_bStaShowingInputWnd)
            {
                return S_OK;
            }
            CTabTipSystemBase::m_bStaShowingInputWnd = false;

            if (m_hTabtipNotifyDbgWnd)
            {
                ::PostMessage(m_hTabtipNotifyDbgWnd, WM_HIDEINPUTDBGWND, 0, 0);
            }
        }
    }
    return S_OK;
}

void UnpackImmersiveFocusTrackingBuffer(PVOID buffer, ULONG BufferSize)
{
    typedef struct _StuBuffer
    {
        union
        {
            HWND hAncestor;
            __int64 noop;
        };
        union
        {
            HWND hFocusWnd;
            __int64 noop;
        };
        union
        {
            char flag;
            DWORD unknow;
        };
        DWORD left;
        int   noop1;
#ifndef __WIN64
        int   noop2;
#endif
        RECT rtFocus;
    }StuBuffer;

    StuBuffer stubuffer;
    int iSize = sizeof(StuBuffer);
    memcpy(&stubuffer, buffer, (sizeof(StuBuffer) <= BufferSize ? sizeof(StuBuffer) : BufferSize));
    assert(iSize == BufferSize);
    bool bTouchMode = stubuffer.flag & 0x40;
    bool bEditFocused = stubuffer.flag & 0x80;

    if (bTouchMode)  // Touch mode
    {
        if (bEditFocused)
        {
            if (stubuffer.rtFocus.left < stubuffer.rtFocus.right && stubuffer.rtFocus.top < stubuffer.rtFocus.bottom)  // rect valid
            {
                CTabTipSystemBase::m_hStaFocusEdit = stubuffer.hFocusWnd;
                CTabTipSystemBase::m_hStaFocusAncestor = stubuffer.hAncestor;
                CTabTipSystemBase::m_rtStaFocusEdit = stubuffer.rtFocus;
                if (CTabTipSystemBase::ms_pIPCWndClient)
                {
                    CTabTipSystemBase::ms_pIPCWndClient->post_message(WM_TABTIP_DISPLAYINPUTWND, SW_SHOW, 0);
                }
                //g_info("[TabTipSystemBase] Touch mode WM_TABTIP_DISPLAYINPUTWND Show\n");
            }
        }
        else
        {
            CTabTipSystemBase::m_hStaFocusEdit = NULL;
            CTabTipSystemBase::m_hStaFocusAncestor = NULL;
            ::SetRectEmpty(&CTabTipSystemBase::m_rtStaFocusEdit);
            if (CTabTipSystemBase::ms_pIPCWndClient)
            {
                CTabTipSystemBase::ms_pIPCWndClient->post_message(WM_TABTIP_DISPLAYINPUTWND, SW_HIDE, 0);
            }
            //g_info("[TabTipSystemBase] Touch mode WM_TABTIP_DISPLAYINPUTWND Hide\n");
        }
    }
    else  // 非触屏
    {
        if (bEditFocused)
        {
            if (stubuffer.hFocusWnd)
            {
                CTabTipSystemBase::m_hStaFocusEdit = stubuffer.hFocusWnd;
                CTabTipSystemBase::m_hStaFocusAncestor = stubuffer.hAncestor;
                CTabTipSystemBase::m_rtStaFocusEdit = stubuffer.rtFocus;
                CTabTipSystemBase::m_bStaStartListenMouseUp = false;
                CTabTipSystemBase::m_hStaFocusEditPointHWnd = WindowFromPoint(POINT{
                    (stubuffer.rtFocus.left + stubuffer.rtFocus.right) >> 1,
                    (stubuffer.rtFocus.top + stubuffer.rtFocus.bottom) >> 1 });
                if (CTabTipSystemBase::ms_pIPCWndClient)
                {
                    CTabTipSystemBase::ms_pIPCWndClient->post_message(WM_TABTIP_DISPLAYINPUTWND, SW_SHOW, 0);
                }
                //g_info("[TabTipSystemBase] Mouse mode WM_TABTIP_DISPLAYINPUTWND Show\n");
            }
        }
        else
        {
            if (stubuffer.hFocusWnd)
            {
                CTabTipSystemBase::m_hStaFocusEdit = NULL;
                CTabTipSystemBase::m_hStaFocusAncestor = NULL;
                ::SetRectEmpty(&CTabTipSystemBase::m_rtStaFocusEdit);
                if (CTabTipSystemBase::ms_pIPCWndClient)
                {
                    CTabTipSystemBase::ms_pIPCWndClient->post_message(WM_TABTIP_DISPLAYINPUTWND, SW_HIDE, 0);
                }
                //g_info("[TabTipSystemBase] Mouse mode WM_TABTIP_DISPLAYINPUTWND Hide\n");
            }
        }
    }
}

NTSTATUS NTAPI WnfCallback(
    WNF_STATE_NAME StateName,
    WNF_CHANGE_STAMP                  ChangeStamp,
    PWNF_TYPE_ID                      TypeId,
    PVOID                             CallbackContext,
    PVOID                             Buffer,
    ULONG                             BufferSize)
{
    if (StateName.Data[0] == WNF_TKBN_IMMERSIVE_FOCUS_TRACKING.Data[0] &&
        StateName.Data[1] == WNF_TKBN_IMMERSIVE_FOCUS_TRACKING.Data[1]
        && Buffer && (LONG)BufferSize > 0)
    {
        UnpackImmersiveFocusTrackingBuffer(Buffer, BufferSize);
    }
    return 0;
}

BOOL CTabTipSystemBase::InstallImmersiveFocusTracking()
{
    typedef NTSTATUS(WINAPI* RtlSubscribeWnfStateChangeNotification) (
        _Outptr_ PWNF_USER_SUBSCRIPTION* Subscription, 
        _In_ WNF_STATE_NAME StateName, 
        _In_ WNF_CHANGE_STAMP ChangeStamp,
        decltype(WnfCallback),
        _In_opt_ PVOID CallbackContext, 
        _In_opt_ PCWNF_TYPE_ID TypeId, 
        _In_opt_ ULONG SerializationGroup, 
        _In_opt_ ULONG Unknown);

    typedef NTSTATUS(WINAPI* RtlQueryWnfStateData)(
        uint32_t* nChangeStamp, 
        WNF_STATE_NAME StateName,
        decltype(WnfCallback),
        size_t* pBuffer, 
        size_t* pBufferSize);

    HINSTANCE hDLL = GetModuleHandle(TEXT("NtDll.dll")) ? GetModuleHandle(TEXT("NtDll.dll")) : LoadLibrary(TEXT("NtDll.dll"));
    if (hDLL)
    {
        RtlSubscribeWnfStateChangeNotification pSubscribe = 
            (RtlSubscribeWnfStateChangeNotification)GetProcAddress(hDLL, "RtlSubscribeWnfStateChangeNotification");
        RtlQueryWnfStateData pQuery = 
            (RtlQueryWnfStateData)GetProcAddress(hDLL, "RtlQueryWnfStateData");
        if (pSubscribe && pQuery)
        {
            uint32_t nChangeStamp{};
            NTSTATUS ntStatus = pQuery(&nChangeStamp, WNF_TKBN_IMMERSIVE_FOCUS_TRACKING, WnfCallback, 0, 0);
            if (ntStatus == 0)
            {
                ntStatus = pSubscribe(&g_pStaUserSubscription,
                    WNF_TKBN_IMMERSIVE_FOCUS_TRACKING,
                    nChangeStamp,
                    WnfCallback, &WNF_TKBN_IMMERSIVE_FOCUS_TRACKING, 0, 0, 0);
                if (ntStatus == 0)
                {
                    g_info("[TabTipSystemBase] RtlSubscribeWnfStateChangeNotification WNF_TKBN_IMMERSIVE_FOCUS_TRACKING succeeded (%08x)\r", nChangeStamp);
                    return TRUE;
                }
                else
                {
                    g_info("[TabTipSystemBase] RtlSubscribeWnfStateChangeNotification WNF_TKBN_IMMERSIVE_FOCUS_TRACKING error : %08x\r", ntStatus);
                }
            }
            else
            {
                g_warning("[TabTipSystemBase] RtlQueryWnfStateData error : %08x\r", ntStatus);
            }
        }
        else
        {
            g_warning("[TabTipSystemBase] Get ntdll interface error,  RtlSubscribeWnfStateChangeNotification: 0x%x RtlQueryWnfStateData:0x%x\r", pSubscribe, pQuery);
        }
    }
    return FALSE;
}

BOOL CTabTipSystemBase::UninstallImmersiveFocusTracking()
{
    if (g_pStaUserSubscription)
    {
        typedef NTSTATUS(WINAPI* RtlUnsubscribeWnfStateChangeNotification)(PWNF_USER_SUBSCRIPTION Subscription);
        RtlUnsubscribeWnfStateChangeNotification pUnsubscriber = (RtlUnsubscribeWnfStateChangeNotification)GetProcAddress(
            LoadLibrary(TEXT("NtDll.dll")), "RtlUnsubscribeWnfStateChangeNotification");
        if (pUnsubscriber)
        {
            pUnsubscriber(g_pStaUserSubscription);
            g_pStaUserSubscription = NULL;
            return TRUE;
        }
    }

    return FALSE;
}

BOOL CTabTipSystemBase::Initialize()
{
    DWORD pid = GetCurrentProcessId();
    DWORD dwSessionId = 0;
    ProcessIdToSessionId(pid, &dwSessionId);
    wsprintf(m_szBufImmersiveFocusEventName, L"%s%d%s", L"Local\\", dwSessionId, ImmersiveFocusEventName);

    SECURITY_DESCRIPTOR secutityDesc;
    ::InitializeSecurityDescriptor(&secutityDesc, SECURITY_DESCRIPTOR_REVISION);
    ::SetSecurityDescriptorDacl(&secutityDesc, TRUE, NULL, FALSE);
    SECURITY_ATTRIBUTES securityAttr;
    securityAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
    securityAttr.bInheritHandle = FALSE;
    securityAttr.lpSecurityDescriptor = &secutityDesc;
    m_hImmersiveEvent = CreateEventW(&securityAttr, TRUE, FALSE, m_szBufImmersiveFocusEventName);
    SetEvent(m_hImmersiveEvent);
    return TRUE;
}

BOOL CTabTipSystemBase::Uninitialize()
{
    if (m_hImmersiveEvent)
    {
        //ResetEvent(m_hImmersiveEvent);  不必要强制重置事件  避免影响TabTip.exe
        CloseHandle(m_hImmersiveEvent);
        m_hImmersiveEvent = NULL;
    }
    return FALSE;
}